/* Copyright (c) 2023 Renmin University of China
RMDB is licensed under Mulan PSL v2.
You can use this software according to the terms and conditions of the Mulan PSL v2.
You may obtain a copy of Mulan PSL v2 at:
        http://license.coscl.org.cn/MulanPSL2
THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
See the Mulan PSL v2 for more details. */

#pragma once
#include "errors.h"
#include "execution_defs.h"
#include "execution_manager.h"
#include "executor_abstract.h"
#include "index/ix.h"
#include "parser/ast.h"
#include "system/sm.h"

class AggregationExecutor : public AbstractExecutor {
   private:
    int curr_idx = -1;
    size_t len_;  
    
    std::vector<Condition> having_conds_;
    std::vector<ColMeta> group_cols_;
    std::vector<ColMeta> sel_cols_;
    std::vector<ColMeta> sel_cols_initial_;
    std::unique_ptr<AbstractExecutor> prev_;
    std::map<std::string, int> aggr_idx_;
    std::vector<std::vector<std::unique_ptr<RmRecord>>> aggr_records_;

    bool empty_table_aggr_ = false;

   public:
    AggregationExecutor(std::unique_ptr<AbstractExecutor> prev, const std::vector<TabCol> &sel_cols, 
        const std::vector<TabCol> &group_cols, const std::vector<Condition> &having_conds)
        : having_conds_(having_conds), prev_(std::move(prev)) {

        const auto &prev_cols = prev_->cols();

        init_sel_cols(sel_cols, prev_cols);
        calculate_offsets();
        init_group_cols(group_cols, prev_cols);
        
    }

    void beginTuple() override {
        for (prev_->beginTuple(); !prev_->is_end(); prev_->nextTuple()) {
            auto record = prev_->Next();
            std::string key = generateGroupKey(record);
            if (aggr_idx_.find(key) == aggr_idx_.end()) {
                aggr_idx_[key] = aggr_records_.size();
                aggr_records_.emplace_back();
            }
            aggr_records_[aggr_idx_[key]].push_back(std::move(record));
        }
        nextTuple();
    }

    void nextTuple() override {
        while (true) {
            if (aggr_records_.empty() && group_cols_.empty()) {
                if (!empty_table_aggr_) {
                    empty_table_aggr_ = true;
                    return;
                }
            }
            ++curr_idx;
            if (is_end()) return;
            updateCurrentRecords();
            if (evalConditions()) break;
        }
    }

    [[nodiscard]] bool is_end() const override{
        return curr_idx >= (int)aggr_records_.size();
    }

    [[nodiscard]] const std::vector<ColMeta> &cols() const override{
        return sel_cols_;
    }

    std::unique_ptr<RmRecord> Next() override {
        std::vector<Value> values;
        values.reserve(sel_cols_initial_.size());

        for (const auto& sel_col : sel_cols_initial_) {
            values.push_back(aggregate_value(sel_col));
        }

        auto data = std::make_unique<char[]>(len_);
        size_t offset = 0;

        for (size_t i = 0; i < sel_cols_.size(); ++i) {
            const auto& col = sel_cols_[i];
            const auto& val = values[i];
            memcpy(data.get() + offset, val.raw->data, val.raw->size);
            offset += col.len;
        }

        return std::make_unique<RmRecord>(len_, data.get());
    }


    Rid &rid() override { return _abstract_rid; }

private:

    std::vector<std::unique_ptr<RmRecord>> curr_records;
    void updateCurrentRecords() {
        curr_records.clear();
        auto &grouped_records = aggr_records_[curr_idx];
        curr_records.reserve(grouped_records.size());
        for (auto &record : grouped_records) {
            curr_records.push_back(std::move(record));
        }
    }

    std::string generateGroupKey(const std::unique_ptr<RmRecord> &record) {
        std::string key;
        for (const auto &group_col : group_cols_) {
            key.append(record->data + group_col.offset, group_col.len);
        }
        return key;
    }

    void init_group_cols(const std::vector<TabCol> &group_cols, const std::vector<ColMeta> &prev_cols) {
        for (const auto &sel_group_col : group_cols) {
            auto pos = get_col(prev_cols, sel_group_col, false);
            group_cols_.push_back(*pos);
        }
    }

    void init_sel_cols(const std::vector<TabCol> &sel_cols, const std::vector<ColMeta> &prev_cols) {
        for (const auto &sel_col : sel_cols) {
            if (sel_col.group_by_cond == ast::COND_COUNT && sel_col.col_name == "*") {
                ColMeta col = ColMeta::createCountAllColMeta(sel_col);
                sel_cols_initial_.push_back(col);
                sel_cols_.push_back(col);
                continue;
            }
            auto pos = get_col(prev_cols, sel_col);
            auto col = *pos;
            col.group_by_cond = sel_col.group_by_cond;
            sel_cols_initial_.push_back(col);
            if (col.group_by_cond != ast::NONE) {
                adjust_col_meta(col, sel_col.group_by_cond);
            }
            sel_cols_.push_back(col);
        }
        calculate_offsets();
    }

    void adjust_col_meta(ColMeta &col, ast::GroupByCond group_by_cond) {
        if (group_by_cond == ast::COND_COUNT) {
            col.type = TYPE_INT;
            col.len = sizeof(int);
        } else if (group_by_cond == ast::COND_SUM) {
            if (col.type == TYPE_STRING) {
                col.type = TYPE_INT;
                col.len = sizeof(int);
            }
        }
    }

    void calculate_offsets() {
        size_t offset = 0;
        for (auto &col : sel_cols_) {
            col.offset = offset;
            offset += col.len;
        }
        len_ = offset;
    }

    bool evalConditions() {
        std::vector<int> indices_to_delete;

        auto should_delete_record = [this](RmRecord* record) -> bool {
            char *base = record->data;
            return !std::all_of(having_conds_.begin(), having_conds_.end(), [base, this](const Condition& cond) {
                return evaluateCondition(cond, base);
            });
        };

        for (int index = 0; index < curr_records.size(); ++index) {
            if (should_delete_record(curr_records[index].get())) {
                indices_to_delete.push_back(index);
            }
        }

        if (indices_to_delete.size() == curr_records.size()) return false;

        for (int i = indices_to_delete.size() - 1; i >= 0; --i) {
            curr_records.erase(curr_records.begin() + indices_to_delete[i]);
        }
    
        return true;
    }

    bool evaluateCondition(const Condition& cond, char* base) {
        if (cond.lhs_col.group_by_cond == ast::NONE) {
            auto value = convertValue(base, *get_col(sel_cols_initial_, cond.lhs_col, true));
            return cond.eval_with_rvalue(value);
        } else {
            ColMeta col_meta = getColMetaForCondition(cond);
            auto value = aggregate_value(col_meta);
            return cond.eval_with_rvalue(value);
        }
    }

    ColMeta getColMetaForCondition(const Condition& cond) {
        if (cond.lhs_col.group_by_cond == ast::COND_COUNT && cond.lhs_col.col_name == "*") {
            return ColMeta::createCountAllColMeta(cond.lhs_col);
        } else {
            return *get_col(sel_cols_initial_, cond.lhs_col, true);
        }
    }

    Value aggregate_value(const ColMeta& sel_col) {
        switch (sel_col.group_by_cond) {
            case ast::NONE:
                return aggregateNone(curr_records, sel_col);
            case ast::COND_COUNT:
                return aggregateCount(curr_records);
            case ast::COND_MAX:
                return aggregateMax(curr_records, sel_col);
            case ast::COND_MIN:
                return aggregateMin(curr_records, sel_col);
            case ast::COND_SUM:
                return aggregateSum(curr_records, sel_col);
        }
    }

    Value aggregateCount(const std::vector<std::unique_ptr<RmRecord>>& curr_records) {
        Value val;
        val.type = TYPE_INT;
        val.set_int(curr_records.size());
        val.init_raw(sizeof(int));
        return val;
    }

    Value aggregateMax(const std::vector<std::unique_ptr<RmRecord>>& curr_records, const ColMeta& sel_col) {
        // if(curr_records.empty())
        //     throw RMDBError("aggre record empty!");
        Value val = convertValue(curr_records[0]->data, sel_col);
        for (const auto& record : curr_records) {
            Value tmp = convertValue(record->data, sel_col);
            if (tmp > val) val = std::move(tmp);
        }
        val.init_raw(sel_col.len);
        return val;
    }

    Value aggregateMin(const std::vector<std::unique_ptr<RmRecord>>& curr_records, const ColMeta& sel_col) {
        // if(curr_records.empty())
        //     throw RMDBError("aggre record empty!");
        Value val = convertValue(curr_records[0]->data, sel_col);
        for (const auto& record : curr_records) {
            Value tmp = convertValue(record->data, sel_col);
            if (tmp < val) val = std::move(tmp);
        }
        val.init_raw(sel_col.len);
        return val;
    }

    Value aggregateSum(const std::vector<std::unique_ptr<RmRecord>>& curr_records, const ColMeta& sel_col) {
        Value val;

        switch (sel_col.type) {
            case TYPE_INT: {
                int sum = sumInt(curr_records, sel_col);
                val.set_int(sum);
                val.init_raw(sizeof(int));
                break;
            }
            case TYPE_FLOAT: {
                float sum = sumFloat(curr_records, sel_col);
                val.set_float(sum);
                val.init_raw(sizeof(float));
                break;
            }
            case TYPE_STRING: {
                return sumString(curr_records, sel_col);
            }
        }

        return val;
    }

    int sumInt(const std::vector<std::unique_ptr<RmRecord>>& curr_records, const ColMeta& sel_col) {
        int sum = 0;
        for (const auto& record : curr_records) {
            sum += *(int*)(record->data + sel_col.offset);
        }
        return sum;
    }

    float sumFloat(const std::vector<std::unique_ptr<RmRecord>>& curr_records, const ColMeta& sel_col) {
        float sum = 0;
        for (const auto& record : curr_records) {
            sum += *(float*)(record->data + sel_col.offset);
        }
        return sum;
    }

    Value sumString(const std::vector<std::unique_ptr<RmRecord>>& curr_records, const ColMeta& sel_col) {
        Value val;
        float sum_float = 0;
        bool is_float = false;

        for (const auto& record : curr_records) {
            const char* base = record->data + sel_col.offset;
            std::string str_value;

            // Extract numerical characters and decimal point
            for (size_t i = 0; i < sel_col.len; ++i) {
                if (std::isdigit(base[i]) || base[i] == '.') {
                    str_value += base[i];
                    if (base[i] == '.') {
                        is_float = true;
                    }
                } else {
                    break;
                }
            }

            // Convert extracted string to float and accumulate
            if (!str_value.empty()) {
                sum_float += std::stof(str_value);
            }
        }

        // Set the result value based on whether sum includes floating point numbers
        if (is_float) {
            val.set_float(sum_float);
            val.type = TYPE_FLOAT;
            val.init_raw(sizeof(float));
        } else {
            val.set_int(static_cast<int>(sum_float));
            val.init_raw(sizeof(int));
        }

        return val;
    }


    Value aggregateNone(const std::vector<std::unique_ptr<RmRecord>>& curr_records, const ColMeta& sel_col) {
        Value val = convertValue(curr_records[0]->data, sel_col);
        val.init_raw(sel_col.len);
        return val;
    }

};